#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include <time.h>

#define PORT 53
#define MAX_LEN 512
#define NAME_SIZE 255
#define NUM_RRS 15
#define NUM_IPADRESS 10
#define MAX_PACKET_SIZE 512

#define SET_UINT16( num, buff) num = htons(*(unsigned short*)*buff); *buff += 2
#define SET_UINT32( num, buff) num = htonl(*(unsigned int*)*buff); *buff += 4

char cname[32] = {0};
struct hostent_nslookup{
    char h_name[32];
    char h_aliases[NUM_RRS][32];
    int h_addrtype;
    int h_length;
    int h_answertype;	/* 1 means auth, 0 means no-auth, 2 means none */
	int h_responetime;
	int h_errocede;
    unsigned int h_addr_list[NUM_RRS];
};

enum{
		ERROR_NONE,
		ERROR_TIMEOUT,
		ERROR_CANNOTRESOLVE
		
	};

struct hostent_nslookup g_nslookup;

typedef struct {
    unsigned short TransId;
    unsigned short Flag;
    unsigned short Question;
    unsigned short Answer;
    unsigned short Authority;
    unsigned short Additional;
    unsigned short Hostaddr;
    unsigned short Class;
} PacketTemp;
/*****************************************************************************/
struct dns_rr{
    char name[NAME_SIZE];
    unsigned short type;
    unsigned short class;
    unsigned short ttl;
    unsigned short rdatalen;
    char data[NAME_SIZE];
};
/*****************************************************************************/
union header_flags {
    unsigned short flags;

    struct {
        unsigned short int rcode:4;
        unsigned short int unused:3;
        unsigned short int recursion_avail:1;
        unsigned short int want_recursion:1;
        unsigned short int truncated:1;
        unsigned short int authorative:1;
        unsigned short int opcode:4;
        unsigned short int question:1;
    } f;
};
/*****************************************************************************/
struct dns_header_s{
    unsigned short id;
    union header_flags flags;
    unsigned short qdcount;
    unsigned short ancount;
    unsigned short nscount;
    unsigned short arcount;
};
/*****************************************************************************/
struct dns_message{
    struct dns_header_s header;
    struct dns_rr question[NUM_RRS];
    struct dns_rr answer[NUM_RRS];
};
/*static char probe_pkt[36] =
  {0x0, 0x0, 0x1, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
  0x1, 'a', 0xc, 'r', 'o', 'o', 't', '-', 's', 'e', 'r', 'v',
  'e', 'r', 's', 0x3, 'n', 'e', 't', 0x0, 0x0, 0x1, 0x0, 0x1, };
  */

static int is_pointer(int in){
	return ((in & 0xc0) == 0xc0);
}

static int DnsSockInit(void) 
{
    struct sockaddr_in sa;
    struct in_addr ip;
    int Dns_QuerySock;
    memset((void *)&sa, 0, sizeof(sa));
    Dns_QuerySock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if ( Dns_QuerySock < 0 ) {
        printf("Could not create query socket");
        return -1;
    }
    ip.s_addr = INADDR_ANY;
    sa.sin_family = AF_INET;
    memcpy((void *)&sa.sin_addr, (void *)&ip, sizeof(struct in_addr));
    sa.sin_port = htons(0);
    /* bind query socket */
    if (bind(Dns_QuerySock, (struct sockaddr *)&sa, sizeof(sa)) < 0){
        printf("dns_init: bind: Could not bind to ephmeral port");
        return -1;
    }
    return Dns_QuerySock;
}
static int WritePacket(int sockfd,const char *name, const char *dns) 
{
    struct sockaddr_in sa;
    unsigned short transactionid;
    PacketTemp PacketTemplate;
    int Cnt,hostnamesize,bufsize;
    char Tmpbuf[255] = {0};
    char *msgbuf = NULL;
    char *Tptr = NULL;


    /* structure socket */
    memset(&PacketTemplate,0,sizeof(PacketTemplate));
    memset(&sa, 0, sizeof(sa));
    sa.sin_family = AF_INET;
    inet_aton(dns, &sa.sin_addr);
    sa.sin_port = ntohs(PORT);
    /* fill packet info */
    srandom(time(NULL));
    transactionid = (uint16_t)random();
    PacketTemplate.TransId = htons(transactionid);
    PacketTemplate.Flag = htons(0x0100);
    PacketTemplate.Question = htons(0x0001);
    PacketTemplate.Answer = htons(0x0000);
    PacketTemplate.Authority = htons(0x0000);
    PacketTemplate.Additional = htons(0x0000);    
    PacketTemplate.Hostaddr = htons(0x0001);
    PacketTemplate.Class = htons(0x0001);
    /* parse host name */
    Tptr = NULL;
    hostnamesize = (strlen(name)+2);
    strncpy(&Tmpbuf[1],name,strlen(name)+1);    
    Cnt = 0;

    for (Cnt = 1;Tmpbuf[Cnt] != 0;Cnt++) {
        if ((Tptr = strchr(&Tmpbuf[Cnt],'.')) != NULL) {
            Tmpbuf[Cnt-1] = Tptr - &Tmpbuf[Cnt];
            Cnt += Tmpbuf[Cnt-1];
        }
        else {
            Tmpbuf[Cnt-1] = strlen(&Tmpbuf[Cnt]); 
            Tmpbuf[hostnamesize] = 0;
            break;
        }
    }
    bufsize = sizeof(unsigned short)*8 + hostnamesize;

    if ((msgbuf = malloc(512)) == NULL) { /* dns packet max 512 bytes */
        printf("malloc failer");
        return -1;
    }
    Tptr = msgbuf;
    /* fill query address,*/
    memcpy(Tptr,&PacketTemplate,12); /* cpoy first 12 bytes */
    Tptr += 12;
    memcpy(Tptr,Tmpbuf,hostnamesize);
    Tptr += hostnamesize;
    memcpy(Tptr,&(PacketTemplate.Hostaddr),4);
    if (sendto(sockfd,msgbuf,bufsize,0,(struct sockaddr *)&sa,sizeof(sa)) < 0) {
        printf("send out socket fail");
		free(msgbuf);
		return -1;
    }
    free(msgbuf);
    return 0;
}

static void parse_dns_name(char *buf_start, char *ptr , char *out , int *len)
{
	int n, flag;
	char *pos = out + (*len);
 
	for(;;)
	{
		flag = (int)ptr[0];
		if(flag == 0)
			break;
		if(is_pointer(flag)){
			n = (int)ptr[1];
			ptr = buf_start + n;
			parse_dns_name(buf_start , ptr , out , len);
			break;
		}else{
			ptr ++;
			memcpy(pos , ptr , flag);	
			pos += flag;
			ptr += flag;
			*len += flag;
			if((int)ptr[0] != 0){
				memcpy(pos , "." , 1);
				pos += 1;
				(*len) += 1;
			}
		}
	}
 
}
 

static void dns_decode_name(char *name, char **buf)
{
    int i, k, len, j;
    unsigned char ulen;

    i = k = 0;
    while( **buf ){
        ulen = *(*buf)++;
        len = ulen;

        for( j = 0; j<len && k<NAME_SIZE; j++)
            name[k++] = *(*buf)++;

        if (k<NAME_SIZE) name[k++] = '.';
    }

    (*buf)++;

    /* end of name string terminated with a 0, not with a dot */
    /* when we reach here, k is at most NAME_SIZE */
    if( k > 0 ){
        name[k-1] = 0x00;
    }else{
        name[0] = 0x00;
    }
}
/*****************************************************************************/
/* Decodes the raw packet in buf to create a rr. Assumes buf points to the 
 * start of a rr. 
 * Note: Question rrs dont have rdatalen or rdata. Set is_question when
 * decoding question rrs, else clear is_question
 */
static void dns_decode_rr(struct dns_rr *rr, char **buf, int is_question,char *header, char *buf_start, struct dns_message *m)
{
	/* if the first two bits the of the name are set, then the message has been
	compressed and so the next byte is an offset from the start of the message
	pointing to the start of the name */
	char *ptr = *buf;
	int len = 0;
	
	if( is_question == 1 )
	{
		if( **buf & 0xC0 ){
	        (*buf)++;
	        header += *(*buf)++;;
	        dns_decode_name( rr->name, &header );
	    }else{
	        /* ordinary decode name */
	        dns_decode_name( rr->name, buf );
	    }
	}
	else
	{
		parse_dns_name(buf_start, ptr, rr->name, &len);
		*buf = *buf + 2;
	}

    SET_UINT16( rr->type, buf );
    SET_UINT16( rr->class, buf);

    if( is_question != 1 ){
        SET_UINT32( rr->ttl, buf );
        SET_UINT16( rr->rdatalen, buf );

        /* BRCM message format wrong. drop it */
        if(((*buf - buf_start) >= MAX_PACKET_SIZE) ||
                (((*buf - buf_start) + rr->rdatalen) >= MAX_PACKET_SIZE) ||
                (rr->rdatalen >= MAX_PACKET_SIZE/2))
        {
            m->header.ancount = 0;
            return;
        }
		if(rr->type == 5)
		{
			len=0;
			memset(cname, 0, sizeof(cname));
			parse_dns_name(buf_start, *buf, cname, &len);
		}
        memcpy( rr->data, *buf, rr->rdatalen );
        *buf += rr->rdatalen;
    }
}

static int dns_decode_message(struct dns_message *m, char **buf)
{
    int i;
    char *header_start = *buf;
    char *buf_start = *buf;

    SET_UINT16( m->header.id, buf );  
    SET_UINT16( m->header.flags.flags, buf );

    SET_UINT16( m->header.qdcount, buf );
    SET_UINT16( m->header.ancount, buf );
    SET_UINT16( m->header.nscount, buf );
    SET_UINT16( m->header.arcount, buf );

    /* decode all the question rrs */
    for( i = 0; i < m->header.qdcount && i < NUM_RRS; i++){
        dns_decode_rr( &m->question[i], buf, 1, header_start, buf_start, m);
    }  
    /* decode all the answer rrs */
    for( i = 0; i < m->header.ancount && i < NUM_RRS; i++){
        dns_decode_rr( &m->answer[i], buf, 0, header_start, buf_start, m);
    }

    return 0;
}

static int ReadAndDecodePacket(int sock, struct hostent_nslookup *nslookup)
{
    char buf[512];
    struct dns_message message;
    struct sockaddr_in sa;
    socklen_t salen;
    salen = sizeof(sa);
    char *Tptr = NULL;
    int numread;
    int answer_num = 0;

    int iplist = 0, allist = 0;
    memset(&message,0,sizeof(struct dns_message));
    if ((numread = recvfrom(sock,buf,MAX_LEN,0,(struct sockaddr *)&sa,&salen)) < 0) {
        printf("Read msg error");
        return -1;
    }
    Tptr = buf;
    dns_decode_message(&message,&Tptr);

    /* the packets has been resolved */
    //printf("\n\n-----id 0x%x flags 0x%x qdcount %d ancount %d nscount %d crcount %d\n\n",\
    //        message.header.id, message.header.flags.flags,\
    //        message.header.qdcount, message.header.ancount, message.header.nscount, message.header.arcount);

	if(message.header.flags.flags & 0x40 && message.header.ancount > 0)	/* --- QR  Opcode AA TC RD RA   Zero   Rcode   --- */
	{								/*     1     4     1  1  1  1    3   	 4         */
		nslookup->h_answertype = 1;	
	}
	else if(message.header.ancount > 0)
	{
		nslookup->h_answertype = 0;	/* No-authoritative answer */
	}
	else	/* message.header.ancount == 0 */
	{
		nslookup->h_answertype = 2;	/* None */
		return -1;
	}
	
    for(answer_num = 0; answer_num < message.header.ancount; answer_num++)
    {
        if (message.answer[answer_num].type == 1)	/* it is Host Address */
        {
            if (message.answer[answer_num].data[0] && answer_num < NUM_IPADRESS) /* nslookup need 10 ip address at most */
            {
				memcpy(&nslookup->h_addr_list[iplist++], message.answer[answer_num].data, 4);
            }
        }
        else if (message.answer[answer_num].type == 5)	/* canonical name for an alias */
        {   
            if (message.answer[answer_num].data[0]) 
            {
                strcpy(nslookup->h_aliases[allist++], message.answer[answer_num].name);	/* alias name */
		strcpy(nslookup->h_name, cname);	/* official name */
            }
        }
    }
    return 0;
}
/* timeout : millseconds */
static struct hostent_nslookup *DnsQuery(const char *name, const char *dns, int dns_timeout) 
{
    int Query_socket;
    fd_set fd;
    int ret;
    struct timeval tv;
	struct timeval tv_start;
	struct timeval tv_end;

	memset(&g_nslookup, 0, sizeof(g_nslookup));
	
    struct hostent_nslookup *nslookup = &g_nslookup;

	if(dns == NULL)
		dns = "8.8.8.8";
	/* initialize socket */
    if ((Query_socket = DnsSockInit()) < 0) {
        printf("Init query socket fail\n");
        return NULL;
    }
	gettimeofday(&tv_start, NULL);
    if ((ret = WritePacket(Query_socket,name, dns)) < 0) {
        printf("write packet fail\n");
        close(Query_socket);
        return NULL;
    }
    tv.tv_sec = dns_timeout / 1000;
	tv.tv_usec = 1000*(dns_timeout % 1000);
       
    FD_SET(Query_socket,&fd);
	if(tv.tv_sec == 0 && tv.tv_usec == 0)
	{
		ret = select( Query_socket+1, &fd, NULL, NULL, NULL );
	}
	else
    	ret = select( Query_socket+1, &fd, NULL, NULL, &tv );

	if (ret <= 0) {
		nslookup->h_errocede=ERROR_TIMEOUT;
        return NULL;
    }
    if (FD_ISSET(Query_socket,&fd)) {
        ret = ReadAndDecodePacket(Query_socket, nslookup);
        }
	gettimeofday(&tv_end, NULL);

    close(Query_socket);
	if(ret != -1)
	{
		nslookup->h_responetime = 1000*(tv_end.tv_sec-tv_start.tv_sec)+(tv_end.tv_usec-tv_start.tv_usec)/1000;
	    return nslookup;
	}
	else
	{
		nslookup->h_errocede=ERROR_CANNOTRESOLVE;
		nslookup->h_responetime = 0;
		return NULL;
	}
}

int main(int argc, char *argv[])
{
	int temp=0;
	int i=0;
	struct in_addr in;
    	struct hostent_nslookup *nslookup;

	if(argc < 2){
		printf("Usage : %s <domain name>\n" , argv[0]);
		exit(-1);
	}

	if(argv[3])
    	nslookup = DnsQuery(argv[1], argv[2], atoi(argv[3]));
	else
		nslookup = DnsQuery(argv[1], argv[2], 0);
	if(nslookup == NULL)
	{
		printf("resolve dns error\n");
		printf("Answer type: %d\n", g_nslookup.h_answertype);
		printf("Response time: %d Ms\n", g_nslookup.h_responetime);
		printf("Error code: %d\n", g_nslookup.h_errocede);
		return -1;
	}

	printf("\n*************************************************\n");

	//printf("Answer type: %d\n", nslookup->h_answertype);
	//printf("Response time: %d Ms\n", nslookup->h_responetime);
	//printf("Error code: %d\n", nslookup->h_errocede);
	printf("Official hostname: %s\n", nslookup->h_name);
#if 1
    for(i=0;nslookup->h_addr_list[i]; i++)
    {
		in.s_addr = nslookup->h_addr_list[i];
        printf("Address %2d: %s\n",  temp++, inet_ntoa(in));
    }
	temp=0;
    for(i=0;nslookup->h_aliases[i] && nslookup->h_aliases[i][0]; i++)
    {
        printf("Alias name %2d: %s\n", temp++, nslookup->h_aliases[i]);
    }
	printf("*************************************************\n");
#endif
    return 0;
}

